import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import sys
import os


def create_radar_chart(df, title, output_path):
    labels = df['trials'].unique()  # The axes will represent the different trials
    num_vars = len(labels)

    # Compute angle for each axis
    angles = np.linspace(0, 2 * np.pi, num_vars, endpoint=False).tolist()
    angles += angles[:1]  # Complete the loop
    # EER = 100-EER
    df['ACC'] = 100 - df['EER'] # acc
    # Setup plot
    fig, ax = plt.subplots(figsize=(20, 16), subplot_kw=dict(polar=True))
    best_acc_in_each_trial = df.groupby('trials')['ACC'].max()
    xtrick_labels = [f"{label}\n(Score: {best_acc_in_each_trial[label]:.2f})" for label in labels]
    plt.xticks(angles[:-1], xtrick_labels, fontsize=15)  # Set trial names as labels


    # Normalize EER values based on the maximum EER for each trial to set dynamic range on each axis
    max_eer_per_trial = df.groupby('trials')['ACC'].max()
    scaled_eer = df.apply(lambda x: x['ACC'] / max_eer_per_trial[x['trials']], axis=1)

    # Define colors for each model
    colors = plt.cm.viridis(np.linspace(0, 1, len(df['model'].unique())))
    for i, (model_name, group) in enumerate(df.groupby('model')):
        # Normalize values by the maximum EER in each trial for consistent scaling
        values = group.apply(lambda x: x['ACC'] / max_eer_per_trial[x['trials']], axis=1).tolist()
        values += values[:1]  # Complete the loop
        # values = 100 - values
        print(values)
        # values = [100 - x * 100 for x in values]
        if len(values) != len(angles):
            print(f"Skipping model {model_name} due to mismatch in data")
            continue
        ax.plot(angles, values, 'o-', linewidth=2, label=model_name, color=colors[i])
        ax.fill(angles, values, alpha=0.25, color=colors[i])

    # ax.set_yticklabels([])
    # ax.set_rticks([])
    # Set Y axis labels to display actual EER values corresponding to each axis segment
    max_value = scaled_eer.max()
    ax.set_ylim(0.75, 1.0)  # Setting the limit
    ax.set_rlim(0.75, 1.0)  # Setting the limit
    ax.set_rticks([0.75, 1.0])  # Set five levels; adjust if needed
    # ax.set_rticks([])  # Set five levels; adjust if needed
    ax.set_yticklabels(["75% of Best Model","Best Model"], fontsize=12)
    # ax.set_yticklabels([])

    plt.title(title, size=18, y=1.1)
    plt.legend(loc='upper right', bbox_to_anchor=(0.1, 0.1), fontsize=15)
    plt.savefig(output_path)
    plt.close()

def plot_data_per_time(csv_path, output_folder):
    df = pd.read_csv(csv_path)
    
    for time in df['time'].unique():
        time_data = df[df['time'] == time]
        output_path = os.path.join(output_folder, f"radar_chart_time_{time}.png")
        create_radar_chart(time_data, f'Radar Chart for Time {time}', output_path)
        print(f"Plot saved: {output_path}")

if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage: python script.py <path_to_csv> <output_folder>")
        sys.exit(1)
    
    csv_file_path = sys.argv[1]
    output_folder = sys.argv[2]

    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    
    plot_data_per_time(csv_file_path, output_folder)

    # Usage 
    # python utils/plot_radar.py /VAF/test/result/zhaosheng_202405_cti_no_fusion/seed_123.csv /VAF/test/result/zhaosheng_202405_cti_no_fusion/images
